package com.zsj.util;

import lombok.extern.slf4j.Slf4j;
import sun.misc.BASE64Decoder;

import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * @Description:
 * @Author:
 * @Date:
 */
@Slf4j
public class TestSign {

    public static final String KEY_ALGORITHM = "RSA";
    public static final String SIGNATURE_ALGORITHM = "SHA256withRSA";

    /**
     * RSA密钥对对象
     */
    public static class RsaKeyPair {
        private final String publicKey;
        private final String privateKey;

        public RsaKeyPair(String publicKey, String privateKey) {
            this.publicKey = publicKey;
            this.privateKey = privateKey;
        }

        public PublicKey getPublicKey() throws Exception {
            byte[] keyBytes = new BASE64Decoder().decodeBuffer(this.publicKey);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            PublicKey publicKey = keyFactory.generatePublic(keySpec);
            return publicKey;
        }

        public PrivateKey getPrivateKey() throws Exception {
            byte[] keyBytes = new BASE64Decoder().decodeBuffer(this.privateKey);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            PrivateKey privateKey = keyFactory.generatePrivate(keySpec);
            return privateKey;
        }
    }

    /**
     * 构建RSA密钥对
     *
     * @return RSA密钥对对象
     * @throws NoSuchAlgorithmException 如果无法生成密钥对，则抛出此异常
     */
    public static RsaKeyPair generateKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();

        RSAPublicKey rsaPublicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) keyPair.getPrivate();

        String publicKeyString = Base64.getEncoder().encodeToString(rsaPublicKey.getEncoded());
        String privateKeyString = Base64.getEncoder().encodeToString(rsaPrivateKey.getEncoded());

        return new RsaKeyPair(publicKeyString, privateKeyString);
    }

    /**
     * 加签并获取token
     *
     * @param requestMap    请求参数映射
     * @param apiPrivateKey RSA私钥
     * @return 签名
     * @throws Exception 如果签名过程中发生错误
     */
    public static String getSign(Map<String, Object> requestMap, PrivateKey apiPrivateKey) throws Exception {
        String message = buildMessage(requestMap);
        log.info("before sign msg:{}.", message);
        String signature = sign(message, apiPrivateKey);
        log.info("authorization sign:[{}].", signature);
        return signature;
    }

    private static String buildMessage(Map<String, Object> requestMap) throws UnsupportedEncodingException {
        Set<String> keySet = requestMap.keySet();
        String[] keyArray = keySet.toArray(new String[0]);
        Arrays.sort(keyArray);
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < keyArray.length; i++) {
            if ("sign".equals(keyArray[i]) || "sign_with_rsa".equals(keyArray[i])) {
                continue;
            }
            // 参数值为空，则不参与签名
            if (String.valueOf(requestMap.get(keyArray[i])).isEmpty()) {
                continue;
            }
            // 格式为：key1=value1&key2=value2
            sb.append(keyArray[i]).append("=").append(URLEncoder.encode(String.valueOf(requestMap.get(keyArray[i])),
                    "UTF-8"));
            if (i < keyArray.length - 1) {
                sb.append("&");
            }
        }
        return sb.toString();
    }

    /**
     * 对message加签
     *
     * @param message    待签名的消息
     * @param privateKey RSA私钥
     * @return 签名后的字符串
     * @throws Exception 如果签名过程中发生错误
     */
    public static String sign(String message, PrivateKey privateKey) throws Exception {
        try {
            Signature sign = Signature.getInstance(SIGNATURE_ALGORITHM);
            sign.initSign(privateKey);
            sign.update(message.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(sign.sign());
        } catch (Exception e) {
            log.error("add sign for msg:[{}] error.", message, e);
            throw new Exception("failed to sign for msg:" + message);
        }
    }

    /**
     * 接口返回验签
     *
     * @param requestMap 请求参数映射
     * @param publicKey  RSA公钥
     * @return 验证是否成功
     * @throws Exception 如果验证过程中发生错误
     */
    public static boolean validate(Map<String, Object> requestMap, PublicKey publicKey) throws Exception {
        String message = buildMessage(requestMap);
        String sign = String.valueOf(requestMap.get("sign_with_rsa"));
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
        signature.initVerify(publicKey);
        signature.update(message.getBytes(StandardCharsets.UTF_8));
        if (signature.verify(Base64.getDecoder().decode(sign))) {
            return true;
        }
        log.error("failed to verify sign:{} for message:{}.", sign, message);
        return false;
    }

    public static void main(String[] args) {
        Map<String, Object> requestMap = new HashMap<>();
        requestMap.put("test", "test123");
        requestMap.put("sign", "test123");
        try {
            RsaKeyPair keys = generateKeyPair();
            System.out.println("公钥------------------");
            System.out.println(keys.publicKey);
            System.out.println("私钥------------------");
            System.out.println(keys.privateKey);
            String sign = getSign(requestMap, keys.getPrivateKey());
            System.out.println("sign------------------");
            System.out.println(sign);
            requestMap.put("sign_with_rsa", sign);
            Boolean isValid = validate(requestMap, keys.getPublicKey());
            System.out.println("validation------------------");
            System.out.println(isValid);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
